package org.cybzacg.encryption.utils;

import org.cybzacg.encryption.model.ValidationResult;
import org.cybzacg.encryption.enums.NationalAlgorithm;

import javax.crypto.*;
import javax.crypto.spec.*;
import java.io.*;
import java.security.*;
import java.util.Arrays;
import java.util.concurrent.atomic.AtomicLong;

/**
 * SM3哈希算法工具类
 * 中国商用密码杂凑算法
 * 
 * @author encryption team
 * @version 1.0
 */
public class SM3Utils {
    
    private static final String ALGORITHM_NAME = "SM3";
    private static final int HASH_SIZE = 32; // SM3输出256位=32字节
    private static final int BLOCK_SIZE = 64; // SM3输入块大小512位=64字节
    
    // 统计信息
    private static final AtomicLong totalHashOperations = new AtomicLong(0);
    private static final AtomicLong totalHMACOperations = new AtomicLong(0);
    
    /**
     * 计算SM3哈希值
     * 
     * @param data 输入数据
     * @return SM3哈希值（32字节）
     * @throws NoSuchAlgorithmException 如果算法不可用
     */
    public static byte[] hash(byte[] data) throws NoSuchAlgorithmException {
        if (data == null) {
            throw new IllegalArgumentException("Input data cannot be null");
        }
        
        try {
            MessageDigest md = MessageDigest.getInstance(ALGORITHM_NAME);
            byte[] hash = md.digest(data);
            totalHashOperations.incrementAndGet();
            return hash;
        } catch (NoSuchAlgorithmException e) {
            throw new NoSuchAlgorithmException("SM3 algorithm not available", e);
        }
    }
    
    /**
     * 计算SM3哈希值（字符串输入）
     * 
     * @param data 输入字符串
     * @return SM3哈希值（32字节）
     * @throws NoSuchAlgorithmException 如果算法不可用
     */
    public static byte[] hash(String data) throws NoSuchAlgorithmException {
        if (data == null) {
            throw new IllegalArgumentException("Input string cannot be null");
        }
        
        return hash(data.getBytes(java.nio.charset.StandardCharsets.UTF_8));
    }
    
    /**
     * 计算SM3哈希值（十六进制字符串）
     * 
     * @param data 输入数据
     * @return SM3哈希值的十六进制字符串（64个字符）
     * @throws NoSuchAlgorithmException 如果算法不可用
     */
    public static String hashToHex(byte[] data) throws NoSuchAlgorithmException {
        byte[] hash = hash(data);
        return bytesToHex(hash);
    }
    
    /**
     * 计算SM3哈希值（十六进制字符串）
     * 
     * @param data 输入字符串
     * @return SM3哈希值的十六进制字符串（64个字符）
     * @throws NoSuchAlgorithmException 如果算法不可用
     */
    public static String hashToHex(String data) throws NoSuchAlgorithmException {
        byte[] hash = hash(data);
        return bytesToHex(hash);
    }
    
    /**
     * 计算SM3 HMAC值
     * 
     * @param data 输入数据
     * @param key 密钥
     * @return SM3 HMAC值（32字节）
     * @throws NoSuchAlgorithmException 如果算法不可用
     * @throws InvalidKeyException 如果密钥无效
     */
    public static byte[] hmac(byte[] data, byte[] key) 
            throws NoSuchAlgorithmException, InvalidKeyException {
        if (data == null) {
            throw new IllegalArgumentException("Input data cannot be null");
        }
        
        if (key == null) {
            throw new IllegalArgumentException("Key cannot be null");
        }
        
        try {
            Mac mac = Mac.getInstance("HmacSM3");
            SecretKeySpec keySpec = new SecretKeySpec(key, "HmacSM3");
            mac.init(keySpec);
            byte[] hmac = mac.doFinal(data);
            totalHMACOperations.incrementAndGet();
            return hmac;
        } catch (NoSuchAlgorithmException e) {
            throw new NoSuchAlgorithmException("HmacSM3 algorithm not available", e);
        } catch (InvalidKeyException e) {
            throw new InvalidKeyException("Invalid key for HmacSM3", e);
        }
    }
    
    /**
     * 计算SM3 HMAC值（十六进制字符串）
     * 
     * @param data 输入数据
     * @param key 密钥
     * @return SM3 HMAC值的十六进制字符串（64个字符）
     * @throws NoSuchAlgorithmException 如果算法不可用
     * @throws InvalidKeyException 如果密钥无效
     */
    public static String hmacToHex(byte[] data, byte[] key) 
            throws NoSuchAlgorithmException, InvalidKeyException {
        byte[] hmac = hmac(data, key);
        return bytesToHex(hmac);
    }
    
    /**
     * 验证SM3哈希值
     * 
     * @param data 原始数据
     * @param expectedHash 期望的哈希值
     * @return 验证结果
     */
    public static ValidationResult verifyHash(byte[] data, byte[] expectedHash) {
        try {
            byte[] actualHash = hash(data);
            boolean isValid = MessageDigest.isEqual(actualHash, expectedHash);
            
            return new ValidationResult.Builder()
                    .valid(isValid)
                    .message(isValid ? "SM3 hash verification passed" : "SM3 hash verification failed")
                    .build();
        } catch (Exception e) {
            return new ValidationResult.Builder()
                    .valid(false)
                    .message("SM3 hash verification error: " + e.getMessage())
                    .build();
        }
    }
    
    /**
     * 验证SM3 HMAC值
     * 
     * @param data 原始数据
     * @param key 密钥
     * @param expectedHMAC 期望的HMAC值
     * @return 验证结果
     */
    public static ValidationResult verifyHMAC(byte[] data, byte[] key, byte[] expectedHMAC) {
        try {
            byte[] actualHMAC = hmac(data, key);
            boolean isValid = MessageDigest.isEqual(actualHMAC, expectedHMAC);
            
            return new ValidationResult.Builder()
                    .valid(isValid)
                    .message(isValid ? "SM3 HMAC verification passed" : "SM3 HMAC verification failed")
                    .build();
        } catch (Exception e) {
            return new ValidationResult.Builder()
                    .valid(false)
                    .message("SM3 HMAC verification error: " + e.getMessage())
                    .build();
        }
    }
    
    /**
     * 分块计算SM3哈希值（适用于大文件）
     * 
     * @param data 输入数据
     * @return SM3哈希值（32字节）
     * @throws NoSuchAlgorithmException 如果算法不可用
     */
    public static byte[] hashLargeData(byte[] data) throws NoSuchAlgorithmException {
        if (data == null) {
            throw new IllegalArgumentException("Input data cannot be null");
        }
        
        try {
            MessageDigest md = MessageDigest.getInstance(ALGORITHM_NAME);
            
            // 分块处理大数据
            int blockSize = 8192; // 8KB块
            for (int i = 0; i < data.length; i += blockSize) {
                int length = Math.min(blockSize, data.length - i);
                md.update(data, i, length);
            }
            
            byte[] hash = md.digest();
            totalHashOperations.incrementAndGet();
            return hash;
        } catch (NoSuchAlgorithmException e) {
            throw new NoSuchAlgorithmException("SM3 algorithm not available", e);
        }
    }
    
    /**
     * 计算文件的SM3哈希值
     * 
     * @param filePath 文件路径
     * @return SM3哈希值（32字节）
     * @throws NoSuchAlgorithmException 如果算法不可用
     * @throws IOException 如果文件读取失败
     */
    public static byte[] hashFile(String filePath) 
            throws NoSuchAlgorithmException, IOException {
        if (filePath == null) {
            throw new IllegalArgumentException("File path cannot be null");
        }
        
        try (InputStream fis = new java.io.FileInputStream(filePath);
             java.io.BufferedInputStream bis = new java.io.BufferedInputStream(fis)) {
            
            MessageDigest md = MessageDigest.getInstance(ALGORITHM_NAME);
            byte[] buffer = new byte[8192];
            int bytesRead;
            
            while ((bytesRead = bis.read(buffer)) != -1) {
                md.update(buffer, 0, bytesRead);
            }
            
            byte[] hash = md.digest();
            totalHashOperations.incrementAndGet();
            return hash;
        } catch (NoSuchAlgorithmException e) {
            throw new NoSuchAlgorithmException("SM3 algorithm not available", e);
        }
    }
    
    /**
     * 字节数组转十六进制字符串
     * 
     * @param bytes 字节数组
     * @return 十六进制字符串
     */
    public static String bytesToHex(byte[] bytes) {
        if (bytes == null) {
            return null;
        }
        
        StringBuilder hex = new StringBuilder(bytes.length * 2);
        for (byte b : bytes) {
            hex.append(String.format("%02x", b));
        }
        return hex.toString();
    }
    
    /**
     * 十六进制字符串转字节数组
     * 
     * @param hex 十六进制字符串
     * @return 字节数组
     */
    public static byte[] hexToBytes(String hex) {
        if (hex == null) {
            return null;
        }
        
        int len = hex.length();
        if (len % 2 != 0) {
            throw new IllegalArgumentException("Hex string must have even length");
        }
        
        byte[] data = new byte[len / 2];
        for (int i = 0; i < len; i += 2) {
            data[i / 2] = (byte) ((Character.digit(hex.charAt(i), 16) << 4) +
                                   Character.digit(hex.charAt(i + 1), 16));
        }
        return data;
    }
    
    /**
     * 生成随机盐值
     * 
     * @param length 盐值长度
     * @return 随机盐值
     */
    public static byte[] generateSalt(int length) {
        if (length <= 0) {
            throw new IllegalArgumentException("Salt length must be positive");
        }
        
        byte[] salt = new byte[length];
        new SecureRandom().nextBytes(salt);
        return salt;
    }
    
    /**
     * 生成密钥派生函数
     * 
     * @param password 密码
     * @param salt 盐值
     * @param iterations 迭代次数
     * @param keyLength 密钥长度
     * @return 派生密钥
     */
    public static byte[] deriveKey(char[] password, byte[] salt, int iterations, int keyLength) {
        if (password == null) {
            throw new IllegalArgumentException("Password cannot be null");
        }
        
        if (salt == null) {
            throw new IllegalArgumentException("Salt cannot be null");
        }
        
        if (iterations <= 0) {
            throw new IllegalArgumentException("Iterations must be positive");
        }
        
        if (keyLength <= 0) {
            throw new IllegalArgumentException("Key length must be positive");
        }
        
        try {
            SecretKeyFactory factory = SecretKeyFactory.getInstance("PBKDF2WithHmacSM3");
            PBEKeySpec spec = new PBEKeySpec(password, salt, iterations, keyLength * 8);
            SecretKey key = factory.generateSecret(spec);
            return key.getEncoded();
        } catch (Exception e) {
            throw new RuntimeException("Key derivation failed: " + e.getMessage(), e);
        }
    }
    
    /**
     * 检查SM3算法是否可用
     * 
     * @return 是否可用
     */
    public static boolean isAvailable() {
        try {
            MessageDigest.getInstance(ALGORITHM_NAME);
            return true;
        } catch (NoSuchAlgorithmException e) {
            return false;
        }
    }
    
    /**
     * 获取算法信息
     * 
     * @return 算法信息
     */
    public static String getAlgorithmInfo() {
        return String.format(
            "SM3 Algorithm Information:\n" +
            "  Algorithm: %s\n" +
            "  Hash Size: %d bytes\n" +
            "  Block Size: %d bytes\n" +
            "  Available: %s\n" +
            "  Total Hash Operations: %d\n" +
            "  Total HMAC Operations: %d",
            ALGORITHM_NAME, HASH_SIZE, BLOCK_SIZE, isAvailable(),
            totalHashOperations.get(), totalHMACOperations.get()
        );
    }
    
    /**
     * 重置统计信息
     */
    public static void resetStatistics() {
        totalHashOperations.set(0);
        totalHMACOperations.set(0);
    }
    
    /**
     * 获取统计信息
     * 
     * @return 统计信息
     */
    public static String getStatistics() {
        return String.format(
            "SM3 Statistics:\n" +
            "  Total Hash Operations: %d\n" +
            "  Total HMAC Operations: %d",
            totalHashOperations.get(), totalHMACOperations.get()
        );
    }
    
    /**
     * SM3哈希计算器
     */
    public static class SM3Hasher {
        
        private final MessageDigest digest;
        private long totalBytes = 0;
        
        /**
         * 构造函数
         * 
         * @throws NoSuchAlgorithmException 如果算法不可用
         */
        public SM3Hasher() throws NoSuchAlgorithmException {
            this.digest = MessageDigest.getInstance(ALGORITHM_NAME);
        }
        
        /**
         * 更新哈希计算
         * 
         * @param data 输入数据
         * @return 当前实例
         */
        public SM3Hasher update(byte[] data) {
            if (data != null) {
                digest.update(data);
                totalBytes += data.length;
            }
            return this;
        }
        
        /**
         * 更新哈希计算
         * 
         * @param data 输入数据
         * @param offset 偏移量
         * @param length 长度
         * @return 当前实例
         */
        public SM3Hasher update(byte[] data, int offset, int length) {
            if (data != null && length > 0) {
                digest.update(data, offset, length);
                totalBytes += length;
            }
            return this;
        }
        
        /**
         * 更新哈希计算
         * 
         * @param data 输入字符串
         * @return 当前实例
         */
        public SM3Hasher update(String data) {
            if (data != null) {
                byte[] bytes = data.getBytes(java.nio.charset.StandardCharsets.UTF_8);
                digest.update(bytes);
                totalBytes += bytes.length;
            }
            return this;
        }
        
        /**
         * 完成哈希计算
         * 
         * @return 哈希值（32字节）
         */
        public byte[] digest() {
            byte[] hash = digest.digest();
            totalHashOperations.incrementAndGet();
            return hash;
        }
        
        /**
         * 重置哈希计算器
         * 
         * @return 当前实例
         */
        public SM3Hasher reset() {
            digest.reset();
            totalBytes = 0;
            return this;
        }
        
        /**
         * 获取已处理的字节数
         * 
         * @return 已处理的字节数
         */
        public long getTotalBytes() {
            return totalBytes;
        }
    }
    
    /**
     * SM3 HMAC计算器
     */
    public static class SM3HMAC {
        
        private final Mac mac;
        private final byte[] originalKey;
        
        /**
         * 构造函数
         * 
         * @param key 密钥
         * @throws NoSuchAlgorithmException 如果算法不可用
         * @throws InvalidKeyException 如果密钥无效
         */
        public SM3HMAC(byte[] key) throws NoSuchAlgorithmException, InvalidKeyException {
            this.mac = Mac.getInstance("HmacSM3");
            SecretKeySpec keySpec = new SecretKeySpec(key, "HmacSM3");
            mac.init(keySpec);
            this.originalKey = key.clone();
        }
        
        /**
         * 更新HMAC计算
         * 
         * @param data 输入数据
         * @return 当前实例
         */
        public SM3HMAC update(byte[] data) {
            if (data != null) {
                mac.update(data);
            }
            return this;
        }
        
        /**
         * 完成HMAC计算
         * 
         * @return HMAC值（32字节）
         */
        public byte[] doFinal() {
            byte[] hmac = mac.doFinal();
            totalHMACOperations.incrementAndGet();
            return hmac;
        }
        
        /**
         * 重置HMAC计算器
         * 
         * @return 当前实例
         */
        public SM3HMAC reset() throws InvalidKeyException {
            mac.init(new SecretKeySpec(originalKey, "HmacSM3"));
            return this;
        }
    }
}
